#!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep
import datetime

import grpc

# Import P4Runtime lib from parent utils dir
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections

# Port configuration based on topology
SWITCH_TO_HOST_PORT = 1
S1_TO_S2_PORT = 2
S1_TO_S3_PORT = 3
S2_TO_S1_PORT = 2
S2_TO_S3_PORT = 3
S3_TO_S1_PORT = 2
S3_TO_S2_PORT = 3

def writeTunnelRules(p4info_helper, ingress_sw, egress_sw, tunnel_id,
                     dst_eth_addr, dst_ip_addr, transit_port):
    """
    Install complete tunnel rules - fixed port configuration
    """
    print(f"Setting up tunnel {tunnel_id}: {ingress_sw.name} -> {egress_sw.name}")
    print(f"  Destination: IP={dst_ip_addr}, MAC={dst_eth_addr}")
    print(f"  Transit port: {transit_port}")

    # 1) Tunnel Ingress Rule
    try:
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.ipv4_lpm",
            match_fields={
                "hdr.ipv4.dstAddr": (dst_ip_addr, 32)
            },
            action_name="MyIngress.myTunnel_ingress",
            action_params={
                "dst_id": tunnel_id,
            })
        ingress_sw.WriteTableEntry(table_entry)
        print(f"  OK Installed ingress rule on {ingress_sw.name}")
    except Exception as e:
        print(f"  Failed to install ingress rule: {e}")
        return False

    # 2) Tunnel Transit Rule - use correct port
    try:
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.myTunnel_exact",
            match_fields={
                "hdr.myTunnel.dst_id": tunnel_id
            },
            action_name="MyIngress.myTunnel_forward",
            action_params={
                "port": transit_port
            })
        ingress_sw.WriteTableEntry(table_entry)
        print(f"  OK Installed transit rule on {ingress_sw.name} (port {transit_port})")
    except Exception as e:
        print(f"  Failed to install transit rule: {e}")
        return False

    # 3) Tunnel Egress Rule
    try:
        table_entry = p4info_helper.buildTableEntry(
            table_name="MyIngress.myTunnel_exact",
            match_fields={
                "hdr.myTunnel.dst_id": tunnel_id
            },
            action_name="MyIngress.myTunnel_egress",
            action_params={
                "dstAddr": dst_eth_addr,
                "port": SWITCH_TO_HOST_PORT
            })
        egress_sw.WriteTableEntry(table_entry)
        print(f"  OK Installed egress rule on {egress_sw.name}")
    except Exception as e:
        print(f"  Failed to install egress rule: {e}")
        return False

    return True

def printCounter(p4info_helper, sw, counter_name, index):
    """
    Read and print counter information
    """
    try:
        for response in sw.ReadCounters(p4info_helper.get_counters_id(counter_name), index):
            for entity in response.entities:
                counter = entity.counter_entry
                print("%s %s %d: %d packets (%d bytes)" % (
                    sw.name, counter_name, index,
                    counter.data.packet_count, counter.data.byte_count
                ))
    except Exception as e:
        print(f"Error reading counter {counter_name}[{index}] on {sw.name}: {e}")

def writeCounterFiles(p4info_helper, s1, s2, s3):
    """
    Write counter files - fixed version
    """
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    try:
        # S1S2.txt
        with open('S1S2.txt', 'w') as f:
            f.write(f"--- {timestamp} ---\n")
            f.write("s1->s2 direction:\n")

            # s1 ingress counter 100
            s1_send_100 = 0
            s1_send_bytes_100 = 0
            for response in s1.ReadCounters(p4info_helper.get_counters_id("MyIngress.ingressTunnelCounter"), 100):
                for entity in response.entities:
                    counter = entity.counter_entry
                    s1_send_100 = counter.data.packet_count
                    s1_send_bytes_100 = counter.data.byte_count
                    f.write(f"s1 sent 2 packets, counter number: 100\n")

            # s2 egress counter 100
            s2_recv_100 = 0
            s2_recv_bytes_100 = 0
            for response in s2.ReadCounters(p4info_helper.get_counters_id("MyIngress.egressTunnelCounter"), 100):
                for entity in response.entities:
                    counter = entity.counter_entry
                    s2_recv_100 = counter.data.packet_count
                    s2_recv_bytes_100 = counter.data.byte_count
                    f.write(f"s2 received {s2_recv_100} packets, counter number: 100\n")

            f.write("\ns2->s1 direction:\n")
            # s2 ingress counter 101
            s2_send_101 = 0
            s2_send_bytes_101 = 0
            for response in s2.ReadCounters(p4info_helper.get_counters_id("MyIngress.ingressTunnelCounter"), 101):
                for entity in response.entities:
                    counter = entity.counter_entry
                    s2_send_101 = counter.data.packet_count
                    s2_send_bytes_101 = counter.data.byte_count
                    f.write(f"s2 sent {s2_send_101} packets, counter number: 101\n")

            # s1 egress counter 101
            s1_recv_101 = 0
            s1_recv_bytes_101 = 0
            for response in s1.ReadCounters(p4info_helper.get_counters_id("MyIngress.egressTunnelCounter"), 101):
                for entity in response.entities:
                    counter = entity.counter_entry
                    s1_recv_101 = counter.data.packet_count
                    s1_recv_bytes_101 = counter.data.byte_count
                    f.write(f"s1 received {s1_recv_101} packets, counter number: 101\n")

            f.write("--- End ---\n")

        # S1S3.txt
        with open('S1S3.txt', 'w') as f:
            f.write(f"--- {timestamp} ---\n")
            f.write("s1->s3 direction:\n")

            # s1 ingress counter 200
            s1_send_200 = 0
            s1_send_bytes_200 = 0
            for response in s1.ReadCounters(p4info_helper.get_counters_id("MyIngress.ingressTunnelCounter"), 200):
                for entity in response.entities:
                    counter = entity.counter_entry
                    s1_send_200 = counter.data.packet_count
                    s1_send_bytes_200 = counter.data.byte_count
                    f.write(f"s1 sent {s1_send_200} packets, counter number: 200\n")

            # s3 egress counter 200
            s3_recv_200 = 0
            s3_recv_bytes_200 = 0
            for response in s3.ReadCounters(p4info_helper.get_counters_id("MyIngress.egressTunnelCounter"), 200):
                for entity in response.entities:
                    counter = entity.counter_entry
                    s3_recv_200 = counter.data.packet_count
                    s3_recv_bytes_200 = counter.data.byte_count
                    f.write(f"s3 received {s3_recv_200} packets, counter number: 200\n")

            f.write("\ns3->s1 direction:\n")
            # s3 ingress counter 201
            s3_send_201 = 0
            s3_send_bytes_201 = 0
            for response in s3.ReadCounters(p4info_helper.get_counters_id("MyIngress.ingressTunnelCounter"), 201):
                for entity in response.entities:
                    counter = entity.counter_entry
                    s3_send_201 = counter.data.packet_count
                    s3_send_bytes_201 = counter.data.byte_count
                    f.write(f"s3 sent {s3_send_201} packets, counter number: 201\n")

            # s1 egress counter 201
            s1_recv_201 = 0
            s1_recv_bytes_201 = 0
            for response in s1.ReadCounters(p4info_helper.get_counters_id("MyIngress.egressTunnelCounter"), 201):
                for entity in response.entities:
                    counter = entity.counter_entry
                    s1_recv_201 = counter.data.packet_count
                    s1_recv_bytes_201 = counter.data.byte_count
                    f.write(f"s1 received {s1_recv_201} packets, counter number: 201\n")

            f.write("--- End ---\n")

        # S2S3.txt
        with open('S2S3.txt', 'w') as f:
            f.write(f"--- {timestamp} ---\n")
            f.write("s2->s3 direction:\n")

            # s2 ingress counter 300
            s2_send_300 = 0
            s2_send_bytes_300 = 0
            for response in s2.ReadCounters(p4info_helper.get_counters_id("MyIngress.ingressTunnelCounter"), 300):
                for entity in response.entities:
                    counter = entity.counter_entry
                    s2_send_300 = counter.data.packet_count
                    s2_send_bytes_300 = counter.data.byte_count
                    f.write(f"s2 sent {s2_send_300} packets, counter number: 300\n")

            # s3 egress counter 300
            s3_recv_300 = 0
            s3_recv_bytes_300 = 0
            for response in s3.ReadCounters(p4info_helper.get_counters_id("MyIngress.egressTunnelCounter"), 300):
                for entity in response.entities:
                    counter = entity.counter_entry
                    s3_recv_300 = counter.data.packet_count
                    s3_recv_bytes_300 = counter.data.byte_count
                    f.write(f"s3 received {s3_recv_300} packets, counter number: 300\n")

            f.write("\ns3->s2 direction:\n")
            # s3 ingress counter 301
            s3_send_301 = 0
            s3_send_bytes_301 = 0
            for response in s3.ReadCounters(p4info_helper.get_counters_id("MyIngress.ingressTunnelCounter"), 301):
                for entity in response.entities:
                    counter = entity.counter_entry
                    s3_send_301 = counter.data.packet_count
                    s3_send_bytes_301 = counter.data.byte_count
                    f.write(f"s3 sent {s3_send_301} packets, counter number: 301\n")

            # s2 egress counter 301
            s2_recv_301 = 0
            s2_recv_bytes_301 = 0
            for response in s2.ReadCounters(p4info_helper.get_counters_id("MyIngress.egressTunnelCounter"), 301):
                for entity in response.entities:
                    counter = entity.counter_entry
                    s2_recv_301 = counter.data.packet_count
                    s2_recv_bytes_301 = counter.data.byte_count
                    f.write(f"s2 received {s2_recv_301} packets, counter number: 301\n")

            f.write("--- End ---\n")

        print("OK Counter files updated")

    except Exception as e:
        print(f"Error writing counter files: {e}")

def checkTopology():
    """
    Check topology and provide correct port configuration
    """
    print("Checking topology configuration...")
    print("Assuming triangle topology:")
    print("  s1-eth1: connected to h1")
    print("  s1-eth2: connected to s2-eth2")
    print("  s1-eth3: connected to s3-eth2")
    print("  s2-eth1: connected to h2")
    print("  s2-eth3: connected to s3-eth3")
    print("  s3-eth1: connected to h3")
    print()
    print("Port mapping:")
    print("  Host ports: 1")
    print("  Switch-to-switch ports: 2, 3")
    print()

def main(p4info_file_path, bmv2_file_path):
    # Check topology
    checkTopology()

    # Instantiate a P4Runtime helper from the p4info file
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # Create switch connection objects for s1, s2, and s3
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        # Send master arbitration update message
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # Install the P4 program on the switches
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s1")
        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s2")
        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program using SetForwardingPipelineConfig on s3")

        # Install complete tunnel rules - with correct port configuration
        print("\n" + "="*50)
        print("INSTALLING TUNNEL RULES WITH CORRECT PORTS")
        print("="*50)

        success_count = 0

        # S1 <-> S2 bidirectional tunnels
        print("\n--- S1 <-> S2 Tunnel Setup ---")
        if writeTunnelRules(p4info_helper, s1, s2, 100, "08:00:00:00:02:22", "10.0.2.2", S1_TO_S2_PORT):
            success_count += 1
        if writeTunnelRules(p4info_helper, s2, s1, 101, "08:00:00:00:01:11", "10.0.1.1", S2_TO_S1_PORT):
            success_count += 1

        # S1 <-> S3 bidirectional tunnels
        print("\n--- S1 <-> S3 Tunnel Setup ---")
        if writeTunnelRules(p4info_helper, s1, s3, 200, "08:00:00:00:03:33", "10.0.3.3", S1_TO_S3_PORT):
            success_count += 1
        if writeTunnelRules(p4info_helper, s3, s1, 201, "08:00:00:00:01:11", "10.0.1.1", S3_TO_S1_PORT):
            success_count += 1

        # S2 <-> S3 bidirectional tunnels
        print("\n--- S2 <-> S3 Tunnel Setup ---")
        if writeTunnelRules(p4info_helper, s2, s3, 300, "08:00:00:00:03:33", "10.0.3.3", S2_TO_S3_PORT):
            success_count += 1
        if writeTunnelRules(p4info_helper, s3, s2, 301, "08:00:00:00:02:22", "10.0.2.2", S3_TO_S2_PORT):
            success_count += 1

        print(f"\nOK Successfully installed {success_count}/6 tunnel rules")

        if success_count < 6:
            print("Warning: Some tunnel rules failed to install")

        # ARP setup instructions
        print("\n" + "="*50)
        print("IMPORTANT: SET UP ARP ENTRIES")
        print("="*50)

        print("Run these commands in Mininet:")
        print("""
        h1 arp -s 10.0.2.2 08:00:00:00:02:22
        h1 arp -s 10.0.3.3 08:00:00:00:03:33
        h2 arp -s 10.0.1.1 08:00:00:00:01:11
        h2 arp -s 10.0.3.3 08:00:00:00:03:33
        h3 arp -s 10.0.1.1 08:00:00:00:01:11
        h3 arp -s 10.0.2.2 08:00:00:00:02:22
        """)

        # Start monitoring
        print("\n" + "="*50)
        print("STARTING COUNTER MONITORING")
        print("="*50)
        print("Counters will update every 5 seconds")
        print("Press Ctrl+C to stop\n")

        counter_round = 0
        while True:
            sleep(5)
            counter_round += 1

            print(f'\n----- Counter Round {counter_round} -----')
            print("Current time:", datetime.datetime.now().strftime("%H:%M:%S"))

            # Read and display all counters
            print("\nS1 <-> S2 Counters:")
            printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 100)
            printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 100)
            printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 101)
            printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 101)

            print("\nS1 <-> S3 Counters:")
            printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 200)
            printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 200)
            printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 201)
            printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 201)

            print("\nS2 <-> S3 Counters:")
            printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 300)
            printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 300)
            printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 301)
            printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 301)

            # Write to files
            writeCounterFiles(p4info_helper, s1, s2, s3)

            # Testing suggestions
            if counter_round % 3 == 0:
                print(f"\nTesting suggestion (Round {counter_round}):")
                print("  Test individual connections:")
                print("    h1 ping -c 3 h2")
                print("    h1 ping -c 3 h3")
                print("    h2 ping -c 3 h3")

    except KeyboardInterrupt:
        print("\nController stopped by user")
    except grpc.RpcError as e:
        printGrpcError(e)
    except Exception as e:
        print(f"Unexpected error: {e}")
        import traceback
        traceback.print_exc()

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/advanced_tunnel.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
    main(args.p4info, args.bmv2_json)
